from __future__ import annotations
import struct
import json
from enum import IntEnum
from dataclasses import dataclass
from src.dto import MessageStruct


class LVColorFormat(IntEnum):
    """LVGL 颜色格式枚举"""
    UNKNOWN = 0
    RAW = 1
    RAW_ALPHA = 2
    L8 = 3
    I1 = 4
    I2 = 5
    I4 = 6
    I8 = 7
    A1 = 8
    A2 = 9
    A4 = 10
    A8 = 11
    RGB565 = 12
    ARGB8565 = 13
    RGB565A8 = 14
    RGB888 = 15
    ARGB8888 = 16
    XRGB8888 = 17


class LVImageFlags(IntEnum):
    """LVGL 图像标志"""
    NONE = 0
    MODIFIABLE = 1
    USER_DATA_1 = 2
    USER_DATA_2 = 4
    USER_DATA_3 = 8
    USER_DATA_4 = 16
    COMPRESSED = 32
    ALPHA = 64


@dataclass
class LVImageHeader:
    """
    LVGL 图像头部结构，对应 C 结构体：
    typedef struct {
        uint32_t cf : 5;          // 颜色格式 (LVColorFormat)
        uint32_t always_zero : 3; // 保留位，总是0
        uint32_t reserved : 2;    // 保留位
        uint32_t w : 11;          // 宽度
        uint32_t h : 11;          // 高度
    } lv_image_header_t;
    
    注意：根据LVGL文档，header应该直接包含cf、w、h字段
    """
    always_zero: int = 0     # 必须是0
    w: int = 0              # 宽度 (最大2047)
    h: int = 0              # 高度 (最大2047)  
    cf: LVColorFormat = LVColorFormat.RGB565  # 颜色格式
    
    def pack(self) -> bytes:
        """打包为32位整数的字节表示"""
        # 验证范围
        if not (0 <= self.cf <= 31):  # 5位
            raise ValueError(f"Color format out of range: {self.cf}")
        if not (0 <= self.w <= 2047):  # 11位
            raise ValueError(f"Width out of range: {self.w}")
        if not (0 <= self.h <= 2047):  # 11位
            raise ValueError(f"Height out of range: {self.h}")
        
        # 按位打包: cf(5) + always_zero(3) + reserved(2) + w(11) + h(11) = 32位
        packed = (
            (self.cf & 0x1F) |           # 5位颜色格式
            ((self.always_zero & 0x7) << 5) |   # 3位零
            ((0 & 0x3) << 8) |               # 2位保留位(固定为0)
            ((self.w & 0x7FF) << 10) |          # 11位宽度
            ((self.h & 0x7FF) << 21)            # 11位高度
        )
        
        # 以小端序返回32位整数
        return struct.pack('<I', packed)
    
    @classmethod
    def unpack(cls, data: bytes) -> LVImageHeader:
        """从字节数据解包"""
        if len(data) < 4:
            raise ValueError("Data too short for image header")
        
        packed = struct.unpack('<I', data[:4])[0]
        
        return cls(
            cf=LVColorFormat(packed & 0x1F),
            always_zero=(packed >> 5) & 0x7,
            w=(packed >> 10) & 0x7FF,
            h=(packed >> 21) & 0x7FF
        )


class LVImageDsc(MessageStruct):
    """
    LVGL 图像描述符结构，对应 C 结构体：
    typedef struct {
        lv_image_header_t header;  // 图像头部信息
        uint32_t data_size;        // 数据大小
        const uint8_t * data;      // 指向图像数据的指针
    } lv_image_dsc_t;
    
    """
    header: LVImageHeader
    data_size: int = 0
    data: bytes = b''
    
    def __post_init__(self):
        # 自动设置数据大小
        if self.data:
            self.data_size = len(self.data)
    
    @classmethod
    def from_image_bytes(
        cls, 
        width: int, 
        height: int, 
        rgb565_data: bytes,
        color_format: LVColorFormat = LVColorFormat.RGB565
    ) -> LVImageDsc:
        """从RGB565数据创建图像描述符"""
        header = LVImageHeader(
            always_zero=0,
            w=width,
            h=height,
            cf=color_format
        )
        
        return cls(
            header=header,
            data_size=len(rgb565_data),
            data=rgb565_data
        )
    
    def pack_binary(self) -> tuple[bytes, bytes]:
        """
        打包为header和payload的元组
        返回: (header_json_bytes, payload_data_bytes)
        - header: JSON格式的头部信息（包含图像描述符信息）
        - payload: 只包含图像数据本身
        """
        # 构建JSON头部信息
        header_info = {
            "w": self.header.w,
            "h": self.header.h,
            "cf": self.header.cf.value,
            "cf_name": self.header.cf.name,
            "data_size": self.data_size,
            "always_zero": self.header.always_zero
        }
        
        # 转换为JSON字节
        header_json = json.dumps(header_info, ensure_ascii=False, separators=(',', ':'))
        header_bytes = header_json.encode('utf-8')
        
        # payload只包含图像数据
        payload_bytes = self.data
        
        return header_bytes, payload_bytes
    
    def pack_c_array(self, array_name: str = "image_data") -> str:
        """
        打包为C语言数组格式的字符串，用于嵌入到C代码中
        """
        data_only = self.data  # 只包含图像数据，不包含头部
        
        # 生成C数组
        hex_data = ', '.join(f'0x{b:02x}' for b in data_only)
        
        c_code = f"""
// Image: {self.header.w}x{self.header.h}, Format: {self.header.cf.name}
static const uint8_t {array_name}_data[] = {{
    {hex_data}
}};

static const lv_image_dsc_t {array_name} = {{
    .header.always_zero = 0,
    .header.w = {self.header.w},
    .header.h = {self.header.h},
    .data_size = {self.data_size},
    .header.cf = {self.header.cf.value},  // {self.header.cf.name}
    .data = {array_name}_data,
}};
""".strip()
        
        return c_code
    
    @classmethod
    def unpack_binary(cls, header_bytes: bytes, payload_bytes: bytes) -> LVImageDsc:
        """从header和payload字节数据解包"""
        try:
            # 解析JSON头部
            header_json = header_bytes.decode('utf-8')
            header_info = json.loads(header_json)
            
            # 创建LVGL头部
            header = LVImageHeader(
                cf=LVColorFormat(header_info['cf']),
                always_zero=header_info.get('always_zero', 0),
                w=header_info['w'],
                h=header_info['h']
            )
            
            # 验证数据大小
            expected_data_size = header_info['data_size']
            if len(payload_bytes) != expected_data_size:
                raise ValueError(f"Payload size mismatch: expected {expected_data_size}, got {len(payload_bytes)}")
            
            return cls(
                header=header,
                data_size=expected_data_size,
                data=payload_bytes
            )
            
        except (json.JSONDecodeError, KeyError, ValueError) as e:
            raise ValueError(f"Failed to unpack binary data: {e}")
    
    @classmethod
    def unpack_legacy_binary(cls, data: bytes) -> LVImageDsc:
        """从传统二进制数据解包（向后兼容）"""
        if len(data) < 8:
            raise ValueError("Data too short for image descriptor")
        
        # 解包头部
        header = LVImageHeader.unpack(data[:4])
        
        # 解包数据大小
        data_size = struct.unpack('<I', data[4:8])[0]
        
        # 解包图像数据
        if len(data) < 8 + data_size:
            raise ValueError(f"Data too short: expected {8 + data_size}, got {len(data)}")
        
        image_data = data[8:8 + data_size]
        
        return cls(
            header=header,
            data_size=data_size,
            data=image_data
        )
    
    def get_pixel_count(self) -> int:
        """获取像素总数"""
        return self.header.w * self.header.h
    
    def get_bytes_per_pixel(self) -> float:
        """根据颜色格式获取每像素字节数"""
        format_sizes = {
            LVColorFormat.RGB565: 2,
            LVColorFormat.RGB888: 3,
            LVColorFormat.ARGB8888: 4,
            LVColorFormat.L8: 1,
            LVColorFormat.A8: 1,
            # 更多格式可以按需添加
        }
        return format_sizes.get(self.header.cf, 1)
    
    def validate(self) -> bool:
        """验证图像数据的一致性"""
        expected_size = int(self.get_pixel_count() * self.get_bytes_per_pixel())
        return self.data_size == len(self.data) and len(self.data) == expected_size

